#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简单的WebSocket服务器测试脚本
用于测试WebSocket代理桥接功能
"""

import asyncio
import websockets
import json
import logging
import argparse
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class WebSocketTestServer:
    """简单的WebSocket测试服务器"""
    
    def __init__(self, host: str = "localhost", port: int = 8765):
        """
        初始化WebSocket测试服务器
        
        Args:
            host: 服务器绑定地址
            port: 服务器端口
        """
        self.host = host
        self.port = port
        self.clients = set()
    
    async def register_client(self, websocket):
        """注册新客户端"""
        self.clients.add(websocket)
        client_addr = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        logger.info(f"新客户端连接: {client_addr}, 当前连接数: {len(self.clients)}")
    
    async def unregister_client(self, websocket):
        """注销客户端"""
        self.clients.discard(websocket)
        client_addr = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
        logger.info(f"客户端断开: {client_addr}, 当前连接数: {len(self.clients)}")
    
    async def broadcast_message(self, message: str, sender_websocket=None):
        """广播消息给所有客户端（除了发送者）"""
        if self.clients:
            # 创建广播任务列表
            broadcast_tasks = []
            for client in self.clients.copy():
                if client != sender_websocket and not client.closed:
                    broadcast_tasks.append(client.send(message))
            
            # 并发发送消息
            if broadcast_tasks:
                await asyncio.gather(*broadcast_tasks, return_exceptions=True)
                logger.info(f"消息已广播给 {len(broadcast_tasks)} 个客户端")
    
    async def handle_client(self, websocket, path):
        """处理客户端连接"""
        await self.register_client(websocket)
        
        try:
            # 发送欢迎消息
            welcome_msg = json.dumps({
                "type": "welcome",
                "message": "欢迎连接到WebSocket测试服务器",
                "timestamp": datetime.now().isoformat(),
                "client_count": len(self.clients)
            }, ensure_ascii=False)
            await websocket.send(welcome_msg)
            
            # 监听客户端消息
            async for message in websocket:
                timestamp = datetime.now().strftime('%H:%M:%S')
                client_addr = f"{websocket.remote_address[0]}:{websocket.remote_address[1]}"
                logger.info(f"[{timestamp}] 收到来自 {client_addr} 的消息: {message}")
                
                try:
                    # 尝试解析JSON消息
                    msg_data = json.loads(message)
                    msg_type = msg_data.get('type', 'unknown')
                    
                    if msg_type == 'echo':
                        # 回显消息
                        echo_response = json.dumps({
                            "type": "echo_response",
                            "original_message": message,
                            "timestamp": datetime.now().isoformat(),
                            "server_message": "这是服务器的回显响应"
                        }, ensure_ascii=False)
                        await websocket.send(echo_response)
                        
                    elif msg_type == 'broadcast':
                        # 广播消息
                        broadcast_msg = json.dumps({
                            "type": "broadcast",
                            "from": client_addr,
                            "message": msg_data.get('message', ''),
                            "timestamp": datetime.now().isoformat()
                        }, ensure_ascii=False)
                        await self.broadcast_message(broadcast_msg, websocket)
                        
                    else:
                        # 默认回显处理
                        response = json.dumps({
                            "type": "response",
                            "received_message": message,
                            "timestamp": datetime.now().isoformat(),
                            "server_note": "消息已收到并处理"
                        }, ensure_ascii=False)
                        await websocket.send(response)
                        
                except json.JSONDecodeError:
                    # 非JSON消息，直接回显
                    echo_msg = f"服务器回显: {message} (时间: {datetime.now().strftime('%H:%M:%S')})"
                    await websocket.send(echo_msg)
                
        except websockets.exceptions.ConnectionClosed:
            logger.info(f"客户端连接正常关闭")
        except Exception as e:
            logger.error(f"处理客户端时出错: {str(e)}")
        finally:
            await self.unregister_client(websocket)
    
    async def send_periodic_messages(self):
        """定期向所有客户端发送消息"""
        counter = 1
        while True:
            try:
                if self.clients:
                    periodic_msg = json.dumps({
                        "type": "server_periodic",
                        "counter": counter,
                        "message": f"这是服务器发送的第 {counter} 条定期消息",
                        "timestamp": datetime.now().isoformat(),
                        "active_clients": len(self.clients)
                    }, ensure_ascii=False)
                    
                    await self.broadcast_message(periodic_msg)
                    counter += 1
                
                await asyncio.sleep(30)  # 每30秒发送一次
                
            except Exception as e:
                logger.error(f"发送定期消息时出错: {str(e)}")
                await asyncio.sleep(30)
    
    async def start_server(self, enable_periodic_messages: bool = True):
        """启动WebSocket服务器"""
        logger.info(f"启动WebSocket测试服务器: ws://{self.host}:{self.port}")
        
        # 启动定期消息任务（如果启用）
        if enable_periodic_messages:
            periodic_task = asyncio.create_task(self.send_periodic_messages())
        
        # 启动WebSocket服务器
        async with websockets.serve(self.handle_client, self.host, self.port):
            logger.info("WebSocket服务器已启动，等待客户端连接...")
            logger.info("按 Ctrl+C 停止服务器")
            
            try:
                await asyncio.Future()  # 永远运行
            except KeyboardInterrupt:
                logger.info("收到中断信号，正在关闭服务器...")
                if enable_periodic_messages:
                    periodic_task.cancel()


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='WebSocket测试服务器')
    parser.add_argument('--host', default='localhost', help='服务器绑定地址')
    parser.add_argument('--port', type=int, default=8765, help='服务器端口')
    parser.add_argument('--no-periodic', action='store_true', 
                       help='禁用定期消息发送')
    
    args = parser.parse_args()
    
    # 创建并启动服务器
    server = WebSocketTestServer(args.host, args.port)
    
    try:
        await server.start_server(enable_periodic_messages=not args.no_periodic)
    except KeyboardInterrupt:
        logger.info("服务器已停止")
    except Exception as e:
        logger.error(f"服务器运行时出错: {str(e)}")


if __name__ == "__main__":
    asyncio.run(main())